
#! 数据获取
"""
    这里在数据获取时已经完成了预处理
    
    
  1. get_img_data ：获取图片数据
  2. get_img_size ：获取一个批次内样本随机缩放的尺寸
  3. make_array  ： 将 list形式的batch数据 转化成多个array构成的tuple
  4. data_loader ：批量读取数据
  5. multithread_loader ： 使用飞桨提供的开启多线程读取数据

"""

import numpy as np
#. 函数式编程而设计，用于增强函数功能的库
import functools
#. 引入飞桨库,处理多线程读取数据
import paddle
#. 引入获取数据文件 
from IData import get_insect_names, get_annotations, get_img_data_from_file
#@ image_augment : 对图片进行随机处理函数
from PData import image_augment

#> 1.获取图片数据 
def get_img_data(record, size=640):
    img, gt_boxes, gt_labels, scales = get_img_data_from_file(record)
    # 随机的乱搞数据
    img, gt_boxes, gt_labels = image_augment(img, gt_boxes, gt_labels, size)
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    mean = np.array(mean).reshape((1, 1, -1))
    std = np.array(std).reshape((1, 1, -1))
    img = (img / 255.0 - mean) / std
    img = img.astype('float32').transpose((2, 0, 1))
    return img, gt_boxes, gt_labels, scales


#! 批量数据读取与加速

#> 2.获取一个批次内样本随机缩放的尺寸
def get_img_size(mode):
    if (mode == 'train') or (mode == 'valid'):
        inds = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
        ii = np.random.choice(inds)
        img_size = 320 + ii * 32
    else:
        img_size = 608
    return img_size


#> 3.将 list形式的batch数据 转化成多个array构成的tuple
def make_array(batch_data):
    img_array = np.array([item[0] for item in batch_data], dtype='float32')
    gt_box_array = np.array([item[1] for item in batch_data], dtype='float32')
    gt_labels_array = np.array([item[2] for item in batch_data], dtype='int32')
    img_scale = np.array([item[3] for item in batch_data], dtype='int32')
    return img_array, gt_box_array, gt_labels_array, img_scale


#> 4.批量读取数据
#? 同一批次内图像的尺寸大小必须是一样的，
#? 不同批次之间的大小是随机的，
#? 由上面定义的get_img_size函数产生
def data_loader(datadir, batch_size=10, mode='train'):
    cname2cid = get_insect_names()
    records = get_annotations(cname2cid, datadir)

    def reader():
        if mode == 'train':
            np.random.shuffle(records)
        batch_data = []
        img_size = get_img_size(mode)
        for record in records:
            #print(record)
            img, gt_bbox, gt_labels, im_shape = get_img_data(record,
                                                             size=img_size)
            batch_data.append((img, gt_bbox, gt_labels, im_shape))
            if len(batch_data) == batch_size:
                yield make_array(batch_data)
                batch_data = []
                img_size = get_img_size(mode)
        if len(batch_data) > 0:
            yield make_array(batch_data)

    return reader


#> 5.使用paddle.reader.xmap_readers实现多线程读取数据
#@ 通过使用飞桨提供的API paddle.reader.xmap_readers可以开启多线程读取数据
#@ 数据已经经过预处理 
def multithread_loader(datadir, batch_size=10, mode='train'):
    cname2cid = get_insect_names()
    records = get_annotations(cname2cid, datadir)
    def reader():
        if mode == 'train':
            np.random.shuffle(records)
        img_size = get_img_size(mode)
        batch_data = []
        for record in records:
            batch_data.append((record, img_size))
            if len(batch_data) == batch_size:
                yield batch_data
                batch_data = []
                img_size = get_img_size(mode)
        if len(batch_data) > 0:
            yield batch_data

    def get_data(samples):
        batch_data = []
        for sample in samples:
            record = sample[0]
            img_size = sample[1]
            img, gt_bbox, gt_labels, im_shape = get_img_data(record,
                                                             size=img_size)
            batch_data.append((img, gt_bbox, gt_labels, im_shape))
        return make_array(batch_data)

    mapper = functools.partial(get_data, )

    return paddle.reader.xmap_readers(mapper, reader, 8, 10)